{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Baseten\n",
    "\n",
    "[Baseten](https://baseten.co) provides all the infrastructure you need to deploy and serve ML models performantly, scalably, and cost-efficiently.\n",
    "\n",
    "This example demonstrates using Langchain with models deployed on Baseten."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "To run this notebook, you'll need a [Baseten account](https://baseten.co) and an [API key](https://docs.baseten.co/settings/api-keys).\n",
    "\n",
    "You'll also need to install the Baseten Python package:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install baseten"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import baseten\n",
    "\n",
    "baseten.login(\"YOUR_API_KEY\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Single model call\n",
    "\n",
    "First, you'll need to deploy a model to Baseten.\n",
    "\n",
    "You can deploy foundation models like WizardLM and Alpaca with one click from the [Baseten model library](https://app.baseten.co/explore/) or if you have your own model, [deploy it with this tutorial](https://docs.baseten.co/deploying-models/deploy).\n",
    "\n",
    "In this example, we'll work with WizardLM. [Deploy WizardLM here](https://app.baseten.co/explore/llama) and follow along with the deployed [model's version ID](https://docs.baseten.co/managing-models/manage)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import Baseten"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model\n",
    "wizardlm = Baseten(model=\"MODEL_VERSION_ID\", verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prompt the model\n",
    "\n",
    "wizardlm(\"What is the difference between a Wizard and a Sorcerer?\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chained model calls\n",
    "\n",
    "We can chain together multiple calls to one or multiple models, which is the whole point of Langchain!\n",
    "\n",
    "This example uses WizardLM to plan a meal with an entree, three sides, and an alcoholic and non-alcoholic beverage pairing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import SimpleSequentialChain\n",
    "from langchain import PromptTemplate, LLMChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the first link in the chain\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=[\"cuisine\"],\n",
    "    template=\"Name a complex entree for a {cuisine} dinner. Respond with just the name of a single dish.\",\n",
    ")\n",
    "\n",
    "link_one = LLMChain(llm=wizardlm, prompt=prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the second link in the chain\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=[\"entree\"],\n",
    "    template=\"What are three sides that would go with {entree}. Respond with only a list of the sides.\",\n",
    ")\n",
    "\n",
    "link_two = LLMChain(llm=wizardlm, prompt=prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the third link in the chain\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    input_variables=[\"sides\"],\n",
    "    template=\"What is one alcoholic and one non-alcoholic beverage that would go well with this list of sides: {sides}. Respond with only the names of the beverages.\",\n",
    ")\n",
    "\n",
    "link_three = LLMChain(llm=wizardlm, prompt=prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the full chain!\n",
    "\n",
    "menu_maker = SimpleSequentialChain(\n",
    "    chains=[link_one, link_two, link_three], verbose=True\n",
    ")\n",
    "menu_maker.run(\"South Indian\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
